package com.gaogzhen.datastructure.tree;

import java.util.*;

public class RBTree<K, V> {

    private RBNode<K, V> root;
    private int size;

    static final class RBNode<K, V> {
        final K key;
        V value;
        boolean red;

        private RBNode<K, V> parent;
        private RBNode<K, V> left;
        private RBNode<K, V> right;

        public RBNode(K key, V value, boolean red, RBNode<K, V> parent, RBNode<K, V> left, RBNode<K, V> right) {
            this.key = key;
            this.value = value;
            this.red = red;
            this.parent = parent;
            this.left = left;
            this.right = right;
        }
    }

    public boolean isEmpty() {
        return size == 0;
    }

    public int size() {
        return size;
    }

    // 左旋
    static <K, V> RBNode<K, V> rotateLeft(RBNode<K, V> root, RBNode<K, V> p) {
        RBNode<K, V> r, pp, rl;
        if (p != null && (r = p.right) != null) {
            if ((rl = p.right = r.left) != null)
                rl.parent = p;
            if ((pp = r.parent = p.parent) == null)
                (root = r).red = false;
            else if (pp.left == p)
                pp.left = r;
            else
                pp.right = r;
            r.left = p;
            p.parent = r;

        }
        return root;
    }

    // 右旋
    static <K, V> RBNode<K, V> rotateRight(RBNode<K, V> root, RBNode<K, V> p) {
        RBNode<K, V> l, pp, lr;
        if (p != null && (l = p.left) != null) {
            if ((lr = p.left = l.right) != null)
                lr.parent = p;
            if ((pp = l.parent = p.parent) == null)
                (root = l).red = false;
            else if (pp.left == p)
                pp.left = l;
            else pp.right = l;
            l.right = p;
            p.parent = l;

        }
        return root;
    }

    // 插入后平衡调整
    static <K, V> RBNode<K, V> balanceInsertion(RBNode<K, V> root, RBNode<K, V> x) {
        x.red = true;
        for (RBNode<K, V> xp, xpp, xppl, xppr; ; ) {
            if ((xp = x.parent) == null) {
                x.red = false;
                return x;
            } else if (!xp.red || (xpp = xp.parent) == null)
                return root;
            if (xp == (xppl = xpp.left)) {
                if ((xppr = xpp.right) != null && xppr.red) {
                    xp.red = xppr.red = false;
                    xpp.red = true;
                    x = xpp;
                } else {
                    if (x == xp.right) {
                        root = rotateLeft(root, x = xp);
                        xpp = (xp = x.parent) == null ? null : xp.parent;
                    }
                    if (xp != null) {
                        xp.red = false;
                        if (xpp != null) {
                            xpp.red = true;
                            root = rotateRight(root, xpp);
                        }
                    }
                }
            } else {
                if (xppl != null && xppl.red) {
                    xppl.red = xp.red = false;
                    xpp.red = true;
                    x = xpp;
                } else {
                    if (x == xp.left) {
                        root = rotateRight(root, x = xp);
                        xpp = (xp = x.parent) == null ? null : xp.parent;
                    }
                    if (xp != null) {
                        xp.red = false;
                        if (xpp != null) {
                            xpp.red = true;
                            root = rotateLeft(root, xpp);
                        }
                    }
                }
            }

        }
    }

    // 查找
    @SuppressWarnings({"rawtypes", "unchecked"})
    public RBNode<K, V> findNode(Object k) {
        if (root == null || k == null) return null;
        RBNode<K, V> p = root;
        while (p != null) {
            if (((Comparable) k).compareTo(p.key) < 0)
                p = p.left;
            else if (((Comparable) k).compareTo(p.key) > 0)
                p = p.right;
            else
                return p;
        }
        return null;
    }

    // 插入
    public V add(K k, V v) {
        if (k == null) return null;
        RBNode<K, V> n = new RBNode<>(k, v, true, null, null, null);
        if (root == null) {
            n.red = false;
            root = n;
            size = 1;
            return v;
        }
        RBNode<K, V> pp = null;
        RBNode<K, V> p = root;
        while (p != null) {
            if (((Comparable) k).compareTo(p.key) < 0) {
                pp = p;
                p = p.left;
            } else if (((Comparable) k).compareTo(p.key) > 0) {
                pp = p;
                p = p.right;
            } else {
                V value = p.value;
                p.value = v;
                return value;
            }
        }
        n.parent = pp;
        if (((Comparable) k).compareTo(pp.key) < 0)
            pp.left = n;
        else
            pp.right = n;
        root = balanceInsertion(root, n);
        size++;
        return v;
    }

    // 代码调整（删除)
    static <K,V> RBNode<K,V> balanceDeletion(RBNode<K,V> root,
                                                       RBNode<K,V> x) {
        for (RBNode<K,V> xp, xpl, xpr;;)  {
            if (x == null || x == root)
                return root;
            else if ((xp = x.parent) == null) {
                x.red = false;
                return x;
            }
            else if (x.red) {
                x.red = false;
                return root;
            }
            else if ((xpl = xp.left) == x) {
                if ((xpr = xp.right) != null && xpr.red) {
                    xpr.red = false;
                    xp.red = true;
                    root = rotateLeft(root, xp);
                    xpr = (xp = x.parent) == null ? null : xp.right;
                }
                if (xpr == null)
                    x = xp;
                else {
                    RBNode<K,V> sl = xpr.left, sr = xpr.right;
                    if ((sr == null || !sr.red) &&
                            (sl == null || !sl.red)) {
                        xpr.red = true;
                        x = xp;
                    }
                    else {
                        if (sr == null || !sr.red) {
                            sl.red = false;
                            xpr.red = true;
                            root = rotateRight(root, xpr);
                            xpr = (xp = x.parent) == null ?
                                    null : xp.right;
                        }
                        if (xpr != null) {
                            xpr.red = xp.red;
                            if ((sr = xpr.right) != null)
                                sr.red = false;
                        }
                        if (xp != null) {
                            xp.red = false;
                            root = rotateLeft(root, xp);
                        }
                        x = root;
                    }
                }
            }
            else { // symmetric
                if (xpl != null && xpl.red) {
                    xpl.red = false;
                    xp.red = true;
                    root = rotateRight(root, xp);
                    xpl = (xp = x.parent) == null ? null : xp.left;
                }
                if (xpl == null)
                    x = xp;
                else {
                    RBNode<K,V> sl = xpl.left, sr = xpl.right;
                    if ((sl == null || !sl.red) &&
                            (sr == null || !sr.red)) {
                        xpl.red = true;
                        x = xp;
                    }
                    else {
                        if (sl == null || !sl.red) {
                            sr.red = false;
                            xpl.red = true;
                            root = rotateLeft(root, xpl);
                            xpl = (xp = x.parent) == null ?
                                    null : xp.left;
                        }
                        if (xpl != null) {
                            xpl.red = xp.red;
                            if ((sl = xpl.left) != null)
                                sl.red = false;
                        }
                        if (xp != null) {
                            xp.red = false;
                            root = rotateRight(root, xp);
                        }
                        x = root;
                    }
                }
            }
        }
    }

    // 删除
    public void removeRBTreeNode(K k) {
        if (root == null) {
            return;
        }
        RBNode<K, V> p = findNode(k);
        if (p == null)
            return;

        RBNode<K,V>  pl = p.left, pr = p.right, replacement;
        if (pl != null && pr != null) {
            RBNode<K,V> s = pr, sl;
            while ((sl = s.left) != null) // find successor
                s = sl;
            boolean c = s.red; s.red = p.red; p.red = c; // swap colors
            RBNode<K,V> sr = s.right;
            RBNode<K,V> pp = p.parent;
            if (s == pr) { // p was s's direct parent
                p.parent = s;
                s.right = p;
            }
            else {
                RBNode<K,V> sp = s.parent;
                if ((p.parent = sp) != null) {
                    if (s == sp.left)
                        sp.left = p;
                    else
                        sp.right = p;
                }
                if ((s.right = pr) != null)
                    pr.parent = s;
            }
            p.left = null;
            if ((p.right = sr) != null)
                sr.parent = p;
            if ((s.left = pl) != null)
                pl.parent = s;
            if ((s.parent = pp) == null)
                root = s;
            else if (p == pp.left)
                pp.left = s;
            else
                pp.right = s;
            if (sr != null)
                replacement = sr;
            else
                replacement = p;
        }
        else if (pl != null)
            replacement = pl;
        else if (pr != null)
            replacement = pr;
        else
            replacement = p;
        if (replacement != p) {
            RBNode<K,V> pp = replacement.parent = p.parent;
            if (pp == null)
                root = replacement;
            else if (p == pp.left)
                pp.left = replacement;
            else
                pp.right = replacement;
            p.left = p.right = p.parent = null;
        }

        RBNode<K,V> r = p.red ? root : balanceDeletion(root, replacement);

        if (replacement == p) {  // detach
            RBNode<K,V> pp = p.parent;
            p.parent = null;
            if (pp != null) {
                if (p == pp.left)
                    pp.left = null;
                else if (p == pp.right)
                    pp.right = null;
            }
        }
        size--;
    }

    public List<RBNode<K,V>> midTraversal() {
        List<RBNode<K,V>> list = new ArrayList<>();

        LinkedList<RBNode<K,V>> l = new LinkedList<>();
        RBNode<K,V> current = root;

        while (current != null || !l.isEmpty()) {
            while (current != null) {
                l.push(current);
                current = current.left;
            }
            if (!l.isEmpty()) {
                current = l.pop();
                list.add(current);
                current = current.right;
            }
        }

        return list;
    }

    private class Itr implements Iterator<RBNode<K,V>> {

        Iterator<RBNode<K,V>> iterator;

        Itr() {
            iterator = midTraversal().iterator();
        }

        @Override
        public boolean hasNext() {
            return iterator.hasNext();
        }

        @Override
        public RBNode<K,V> next() {
            return iterator.next();
        }
    }

    @Override
    public String toString() {
        Iterator<RBNode<K,V>> iterator = new Itr().iterator;
        if (!iterator.hasNext())
            return "[]";
        StringBuilder b = new StringBuilder("[");
        while (true) {
            RBNode<K,V> next = iterator.next();
            b.append("(").append(next.key).append(",")
                .append(next.value).append(",")
                .append(next.red? "红": "黑").append(")");
            if (!iterator.hasNext()) {
                return b.append("]").toString();
            }
            b.append(",");
        }
    }

}
